from face_matcher import match_faces, match_multiple_faces
import os

def main():
    # 源图片路径
    source_image = "static/images/src-face/zhangjingyi/zhangjingyi-01.jpg"
    # 源图片目录路径
    source_directory = "static/images/src-face/zhangjingyi"
    # 目标目录路径
    target_directory = "static/images/tar-face/zhangjingyi"
    # 模型存储目录
    # model_directory = "./custom_models"
    
    # 获取目标目录中的所有图片数量
    target_images = [f for f in os.listdir(target_directory) 
                    if f.lower().endswith(('.jpg', '.jpeg', '.png'))]
    total_target_images = len(target_images)
    
    print("=" * 50)
    print("单图匹配模式")
    print("=" * 50)
    print(f"开始匹配源图片 {source_image} 与目标目录 {target_directory} 中的图片...")
    
    # 调用单图匹配函数
    single_matched_images = match_faces(
        source_image, 
        target_directory
        # model_dir=model_directory
    )
    
    # 输出单图匹配结果
    if single_matched_images:
        print("\n匹配成功的图片:")
        for img in single_matched_images:
            print(f"- {img}")
        single_match_count = len(single_matched_images)
        single_match_rate = (single_match_count / total_target_images) * 100
        print(f"共找到 {single_match_count} 张匹配图片，匹配率: {single_match_rate:.2f}%")
    else:
        print("未找到匹配的图片")
        single_match_count = 0
        single_match_rate = 0
    
    print("\n" + "=" * 50)
    print("多图匹配模式")
    print("=" * 50)
    print(f"开始匹配源目录 {source_directory} 中的图片与目标目录 {target_directory} 中的图片...")
    
    # 调用多源图片匹配函数
    match_results = match_multiple_faces(
        source_directory, 
        target_directory
    )
    
    # 输出多图匹配结果
    total_matched = 0
    all_matched_images = set()
    
    print("\n匹配结果汇总:")
    for source_img, matched_images in match_results.items():
        print(f"\n源图片 {source_img} 匹配成功的图片:")
        if matched_images:
            for img in matched_images:
                print(f"- {img}")
                all_matched_images.add(img)
            print(f"共找到 {len(matched_images)} 张匹配图片")
            total_matched += len(matched_images)
        else:
            print("未找到匹配的图片")
    
    multi_match_count = len(all_matched_images)
    multi_match_rate = (multi_match_count / total_target_images) * 100
    print(f"\n多图匹配总计匹配成功 {multi_match_count} 张图片，匹配率: {multi_match_rate:.2f}%")
    
    # 对比匹配率
    print("\n" + "=" * 50)
    print("匹配率对比")
    print("=" * 50)
    print(f"单图匹配: {single_match_count}/{total_target_images} = {single_match_rate:.2f}%")
    print(f"多图匹配: {multi_match_count}/{total_target_images} = {multi_match_rate:.2f}%")
    
    if multi_match_rate > single_match_rate:
        improvement = multi_match_rate - single_match_rate
        print(f"多图匹配提高了匹配率: +{improvement:.2f}%")
    elif multi_match_rate < single_match_rate:
        decrease = single_match_rate - multi_match_rate
        print(f"多图匹配降低了匹配率: -{decrease:.2f}%")
    else:
        print("两种方法匹配率相同")

if __name__ == "__main__":
    main()
